import { Injectable, Logger, OnModuleInit, OnModuleDestroy } from '@nestjs/common';
import { ConfigService } from '@nestjs/config';
import * as WebSocket from 'ws';
import axios from 'axios';
import { EventEmitter } from 'events';
import { Workflow, WorkflowNode, WorkflowParams } from './interfaces/workflow.interface';
import { WorkflowValidator } from './validators/workflow.validator';
import { CacheService } from './services/cache.service';
import { ComfyUIConnectionError, ComfyUIExecutionError } from './interfaces/errors.interface';

interface ComfyUIProgress {
    value: number;
    max: number;
    node: string | null;
}

interface ComfyUIOutput {
    images?: { filename: string; subfolder: string; type: string }[];
    error?: string;
}

@Injectable()
export class ComfyUIService implements OnModuleInit, OnModuleDestroy {
    private readonly logger = new Logger(ComfyUIService.name);
    private readonly wsUrl: string;
    private readonly apiUrl: string;
    private ws: WebSocket | null = null;
    private readonly events = new EventEmitter();
    private readonly taskQueue: Map<string, {
        resolve: (value: any) => void;
        reject: (reason?: any) => void;
        timeout: NodeJS.Timeout;
    }> = new Map();
    private isConnected = false;
    private reconnectAttempts = 0;
    private readonly maxReconnectAttempts = 5;
    private readonly reconnectInterval = 5000;

    constructor(
        private configService: ConfigService,
        private cacheService: CacheService
    ) {
        this.wsUrl = this.configService.get<string>('COMFYUI_WS_URL')!;
        this.apiUrl = this.configService.get<string>('COMFYUI_API_URL')!;
    }

    async onModuleInit() {
        await this.connect();
    }

    onModuleDestroy() {
        this.disconnect();
    }

    private async connect() {
        if (this.isConnected || this.reconnectAttempts >= this.maxReconnectAttempts) {
            return;
        }

        try {
            this.ws = new WebSocket(this.wsUrl);

            this.ws.on('open', () => {
                this.logger.log('Connected to ComfyUI WebSocket');
                this.isConnected = true;
                this.reconnectAttempts = 0;
            });

            this.ws.on('message', (data: WebSocket.Data) => {
                try {
                    const message = JSON.parse(data.toString());
                    this.handleMessage(message);
                } catch (error) {
                    this.logger.error('Failed to parse WebSocket message:', error);
                }
            });

            this.ws.on('close', () => {
                this.logger.warn('ComfyUI WebSocket connection closed');
                this.handleDisconnect();
            });

            this.ws.on('error', (error) => {
                this.logger.error('ComfyUI WebSocket error:', error);
                this.handleDisconnect();
            });

        } catch (error) {
            this.logger.error('Failed to connect to ComfyUI:', error);
            this.handleDisconnect();
        }
    }

    private handleDisconnect() {
        this.isConnected = false;
        this.ws = null;

        if (this.reconnectAttempts < this.maxReconnectAttempts) {
            this.reconnectAttempts++;
            this.logger.log(`Attempting to reconnect (${this.reconnectAttempts}/${this.maxReconnectAttempts})...`);
            setTimeout(() => this.connect(), this.reconnectInterval);
        } else {
            this.logger.error('Max reconnection attempts reached');
            // 清理所有等待中的任务
            this.taskQueue.forEach(({ reject, timeout }) => {
                clearTimeout(timeout);
                reject(new Error('Connection lost'));
            });
            this.taskQueue.clear();
        }
    }

    private handleMessage(message: any) {
        const { type, data } = message;

        switch (type) {
            case 'execution_start':
                this.handleExecutionStart(data);
                break;
            case 'execution_cached':
                this.handleExecutionCached(data);
                break;
            case 'executing':
                this.handleExecuting(data);
                break;
            case 'executed':
                this.handleExecuted(data);
                break;
            case 'execution_error':
                this.handleExecutionError(data);
                break;
            default:
                this.logger.debug(`Unhandled message type: ${type}`);
        }
    }

    private handleExecutionStart(data: { prompt_id: string }) {
        const task = this.taskQueue.get(data.prompt_id);
        if (task) {
            this.events.emit('progress', { value: 0, max: 100, promptId: data.prompt_id });
        }
    }

    private handleExecuting(data: { node: string; prompt_id: string; value: number; max: number }) {
        const progress = (data.value / data.max) * 100;
        this.events.emit('progress', {
            value: progress,
            max: 100,
            node: data.node,
            promptId: data.prompt_id
        });
    }

    private handleExecuted(data: { prompt_id: string; output: ComfyUIOutput }) {
        const task = this.taskQueue.get(data.prompt_id);
        if (task) {
            const { resolve, timeout } = task;
            clearTimeout(timeout);
            this.taskQueue.delete(data.prompt_id);

            if (data.output.images?.length) {
                const imageUrl = this.buildImageUrl(data.output.images[0]);
                resolve(imageUrl);
            } else {
                resolve(null);
            }
        }
    }

    private handleExecutionCached(data: { prompt_id: string; output: ComfyUIOutput }) {
        this.handleExecuted(data);
    }

    private handleExecutionError(data: { prompt_id: string; error: string }) {
        const task = this.taskQueue.get(data.prompt_id);
        if (task) {
            const { reject, timeout } = task;
            clearTimeout(timeout);
            this.taskQueue.delete(data.prompt_id);
            reject(new Error(data.error));
        }
    }

    private buildImageUrl(image: { filename: string; subfolder: string }): string {
        const baseUrl = this.apiUrl.replace('/ws', '');
        const subfolder = image.subfolder ? `${image.subfolder}/` : '';
        return `${baseUrl}/view?filename=${subfolder}${image.filename}`;
    }

    async generateImage(workflow: Workflow, params: WorkflowParams): Promise<string> {
        // 验证参数和工作流
        WorkflowValidator.validateParams(params);
        WorkflowValidator.validateWorkflow(workflow);

        // 检查缓存
        const cachedResult = await this.cacheService.getCachedResult(workflow, params);
        if (cachedResult) {
            return cachedResult;
        }

        if (!this.isConnected) {
            throw new ComfyUIConnectionError('Not connected to ComfyUI');
        }

        try {
            const updatedWorkflow = this.updateWorkflowParams(workflow, params);
            const { prompt_id } = await this.executeWorkflow(updatedWorkflow);

            const result = await new Promise<string>((resolve, reject) => {
                const timeout = setTimeout(() => {
                    this.taskQueue.delete(prompt_id);
                    reject(new ComfyUIExecutionError('Generation timeout'));
                }, 300000);

                this.taskQueue.set(prompt_id, { resolve, reject, timeout });
            });

            // 缓存结果
            await this.cacheService.cacheResult(workflow, params, result);
            return result;

        } catch (error) {
            this.logger.error('Failed to generate image:', error);
            if (error instanceof ComfyUIExecutionError) {
                throw error;
            }
            throw new ComfyUIExecutionError(error.message);
        }
    }

    private updateWorkflowParams(workflow: Workflow, params: WorkflowParams): Workflow {
        const updatedWorkflow = JSON.parse(JSON.stringify(workflow)) as Workflow;

        Object.values(updatedWorkflow.nodes).forEach((node: WorkflowNode) => {
            // 处理 CLIPTextEncode 节点
            if (node.type === 'CLIPTextEncode') {
                if (node.inputs?.text && !node.inputs.text.includes('negative')) {
                    node.inputs.text = params.prompt;
                }
                if (node.inputs?.text?.includes('negative')) {
                    node.inputs.text = params.negative_prompt || '';
                }
            }

            // 处理 KSampler 节点
            if (node.type === 'KSampler') {
                if (params.steps) node.inputs.steps = params.steps;
                if (params.cfg_scale) node.inputs.cfg = params.cfg_scale;
                if (params.seed) node.inputs.seed = params.seed;
            }

            // 处理 EmptyLatentImage 节点
            if (node.type === 'EmptyLatentImage') {
                if (params.width) node.inputs.width = params.width;
                if (params.height) node.inputs.height = params.height;
            }
        });

        return updatedWorkflow;
    }

    private async executeWorkflow(workflow: Workflow) {
        try {
            const response = await axios.post(`${this.apiUrl}/prompt`, workflow);
            return response.data;
        } catch (error) {
            this.logger.error('Failed to execute workflow:', error);
            throw error;
        }
    }

    onProgress(callback: (progress: ComfyUIProgress) => void) {
        this.events.on('progress', callback);
        return () => this.events.off('progress', callback);
    }

    disconnect() {
        if (this.ws) {
            this.ws.close();
            this.ws = null;
            this.isConnected = false;
        }
    }
}